package rb;

import avl.AvlTree;

import java.io.PrintWriter;

public class RBTree {
    static class RBTreeNode {
        public int val;
        public RBTreeNode parent;
        public RBTreeNode left;
        public RBTreeNode right;
        public Color color = Color.Red;
        public RBTreeNode(int val) {
            this.val = val;
        }
    }
    public RBTreeNode root;
    public void insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if(root == null) {
            root = node;
            root.color = Color.Black;
            return;
        }
        RBTreeNode cur = root;
        RBTreeNode parent = root;
        while(cur != null) {
            parent = cur;
            if(cur.val > val) {
                cur = cur.left;
            }else if(cur.val < val) {
                cur = cur.right;
            }else {
                return;
            }
        }
        if(parent.val > val) {
            parent.left = node;
        }else if(parent.val < val){
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        //进行红黑节点的调整
        while(parent != null && parent.color == Color.Red) {
            RBTreeNode grandfather = parent.parent; //必不为null
            if(grandfather.left == parent) {
                RBTreeNode uncle = grandfather.right;
                //情况一
                if(uncle != null && uncle.color == Color.Red) {
                    parent.color = Color.Black;
                    grandfather.color = Color.Red;
                    uncle.color = Color.Black;
                    //判断是否向上还有红色节点
                    cur = grandfather;
                    parent = cur.parent;
                }else {
                    //情况三处理
                    //1.进行左旋
                    if(parent.right == cur) {
                        rotateLeft(parent);
                        //2.交换引用
                        RBTreeNode tmp = cur;
                        cur = parent;
                        parent = tmp;
                    }
                    //处理完毕，此时就为情况二
                    rotateRight(grandfather);
                    grandfather.color = Color.Red;
                    parent.color = Color.Black;
                }
            }else {
                RBTreeNode uncle = grandfather.left;
                if(uncle != null && uncle.color == Color.Red) {
                    parent.color = Color.Black;
                    grandfather.color = Color.Red;
                    uncle.color = Color.Black;
                    //判断是否向上还有红色节点
                    cur = grandfather;
                    parent = cur.parent;
                }else {
                    //情况三处理
                    //1.进行右旋
                    if(parent.left == cur) {
                        rotateRight(parent);
                        //2.交换引用
                        RBTreeNode tmp = cur;
                        cur = parent;
                        parent = tmp;
                    }
                    //处理完毕，此时就为情况二
                    rotateLeft(grandfather);
                    grandfather.color = Color.Red;
                    parent.color = Color.Black;
                }
            }
        }
        root.color = Color.Black;
    }
    public boolean isRBTree(RBTreeNode root) {
        if(root == null) {
            System.out.println("当前树是一棵空树");
            return false;
        }
        if(root.color == Color.Red) {
            System.out.println("违背性质: 根节点不能为红色");
            return false;
        }

        RBTreeNode cur = root;
        int blackCount = 0;
        while(cur != null) {
            if(cur.color == Color.Black) {
                blackCount++;
            }
            cur = cur.left;
        }
        //
        return checkRedColor(root) && checkBlackNum(root,0,blackCount);
    }

    private boolean checkRedColor(RBTreeNode root) {
        if(root == null) {
            return true;
        }
        //只有为红色的时候才会进行判断，黑色不用判断
        if(root.color == Color.Red) {
            //是红色时一定右父亲节点
            RBTreeNode parent = root.parent;
            if(parent.color == Color.Red) {
                System.out.println("违反性质: 连续两个节点不允许为红节点");
                return false;
            }
        }
        return checkRedColor(root.right) && checkRedColor(root.left);
    }

    private boolean checkBlackNum(RBTreeNode root,int pathCount,int blackCount) {
        if(root == null) {
            return true;
        }
        if(root.color == Color.Black) {
            pathCount++;
        }
        //在根节点的时候进行判断
        if(root.right == null && root.left == null) {
            if (pathCount != blackCount) {
                System.out.println("违反性质: 每条路径上的黑色节点个数相同");
                return false;
            }
        }
        return checkBlackNum(root.left,pathCount,blackCount)
                && checkBlackNum(root.right,pathCount,blackCount);
    }

    private void rotateLeft(RBTreeNode parent) {
        RBTreeNode subR = parent.right;
        RBTreeNode subRl = subR.left;
        subR.left = parent;
        RBTreeNode pParent = parent.parent;
        parent.parent = subR;
        parent.right = subRl;
        if(subRl != null) {
            subRl.parent = parent;
        }
        if(root == parent) {
            root = subR;
            root.parent = null;
        }else {
            if(parent == pParent.left) {
                pParent.left = subR;
            }else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }

    }


    private void rotateRight(RBTreeNode parent) {
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;

        parent.left = subLR;
        subL.right = parent;
        RBTreeNode pParent = parent.parent;

        parent.parent = subL;
        if(subLR != null) {
            subLR.parent = parent;
        }

        if(root == parent) {
            //更新当前新的根节点
            root = subL;
            root.parent = null;
        }else {
            //不是根节点，判断当前pParent节点是左子树还是右子树
            if(parent == pParent.right) {
                pParent.right = subL;
            }else {
                pParent.left = subL;
            }
            subL.parent = pParent;
        }
    }
}
